//------------------------------------------------------------------------------
// RegNamespace.cpp
//------------------------------------------------------------------------------
// File provided for Microsoft Shared Source.
// Copyright (c) Microsoft Corporation. All rights reserved.
//
// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF 
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO 
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A 
// PARTICULAR PURPOSE.
//------------------------------------------------------------------------------
//
// Registry namespace module/registration
//
//------------------------------------------------------------------------------

#include "pch.h"
#include <msidefs.h>
#include <RegNamespace.h>
#include <XMLDocument.h>

#define _MERGE_PROXYSTUB

#ifdef _MERGE_PROXYSTUB

extern "C" 
{
    BOOL WINAPI PrxDllMain(__in HINSTANCE hInstance, DWORD dwReason, __in void *pvReserved);
    STDAPI PrxDllCanUnloadNow(void);
    STDAPI PrxDllGetClassObject(REFCLSID rclsid, REFIID riid, __deref_out void **ppv);
    STDAPI PrxDllRegisterServer(void);
    STDAPI PrxDllUnregisterServer(void);
}

#endif // _MERGE_PROXYSTUB

class CRegNamespaceModule : public CAtlDllModuleT<CRegNamespaceModule>
{
public:
    HRESULT RegisterServer(BOOL fTypeLib);
    HRESULT UnregisterServer(BOOL fTypeLib);
    HRESULT RegisterPropertySchema();
    HRESULT UnregisterPropertySchema();
    void NotifyDesktopChange();

private:
    HRESULT _BuildSchemaPath(__out_ecount(cch) PWSTR pszPath, DWORD cch);
};

CRegNamespaceModule _AtlModule;

void CRegNamespaceModule::NotifyDesktopChange()
{
    CComMemPtr<ITEMIDLIST_ABSOLUTE> spidl;
    if (SUCCEEDED(SHGetKnownFolderIDList(FOLDERID_Desktop, KF_FLAG_NO_ALIAS, NULL, &spidl)))
    {
        SHChangeNotify(SHCNE_UPDATEDIR, SHCNF_IDLIST, spidl.m_pData, NULL);
    }
}

HRESULT CRegNamespaceModule::RegisterServer(BOOL fTypeLib)
{
    HRESULT hr = CAtlDllModuleT::RegisterServer(fTypeLib);
#ifdef _MERGE_PROXYSTUB
    if (SUCCEEDED(hr))
    {
        hr = PrxDllRegisterServer();
    }
#endif
    if (SUCCEEDED(hr))
    {
        hr = RegisterPropertySchema();
        if (SUCCEEDED(hr))
        {
            NotifyDesktopChange();
        }
    }
	return hr;
}

HRESULT CRegNamespaceModule::UnregisterServer(BOOL fTypeLib)
{
    HRESULT hr = UnregisterPropertySchema();
    if (SUCCEEDED(hr))
    {
        hr = CAtlDllModuleT::UnregisterServer(fTypeLib);
#ifdef _MERGE_PROXYSTUB
        if (SUCCEEDED(hr))
        {
            hr = PrxDllRegisterServer();
            if (SUCCEEDED(hr))
            {
                hr = PrxDllUnregisterServer();
            }
        }
#endif
        if (SUCCEEDED(hr))
        {
            NotifyDesktopChange();
        }
    }
	return hr;
}

HRESULT CRegNamespaceModule::_BuildSchemaPath(__out_ecount(cch) PWSTR pszPath, DWORD cch)
{
    ATLASSERT(cch >= MAX_PATH);
    CPathString strModule;
    HRESULT hr = GetModuleFileName(_AtlBaseModule.GetModuleInstance(), strModule, strModule.Size()) ? S_OK : AtlHresultFromLastError();
    if (SUCCEEDED(hr))
    {
        hr = StringCchCopy(pszPath, cch, strModule);
        if (SUCCEEDED(hr))
        {
            PathRemoveExtension(pszPath);
            hr = PathAddExtension(pszPath, L".propdesc") ? S_OK : E_FAIL;
        }
    }
    return hr;
}

HRESULT CRegNamespaceModule::RegisterPropertySchema()
{
    CPathString strModule;
    HRESULT hr = GetModuleFileName(_AtlBaseModule.GetModuleInstance(), strModule, strModule.Size()) ? S_OK : AtlHresultFromLastError();
    if (SUCCEEDED(hr))
    {
        CXMLDocument xml;
        hr = xml.Load(_AtlBaseModule.GetModuleInstance(), IDR_REGPROPSCHEMA);
        if (SUCCEEDED(hr))
        {
            CComPtr<IXSLTemplate> spxslTemplate;
            hr = spxslTemplate.CoCreateInstance(__uuidof(XSLTemplate60));
            if (SUCCEEDED(hr))
            {
                hr = spxslTemplate->putref_stylesheet(xml);
                if (SUCCEEDED(hr))
                {
                    CComPtr<IXSLProcessor> spxslProc;
                    hr = spxslTemplate->createProcessor(&spxslProc);
                    if (SUCCEEDED(hr))
                    {
                        CComBSTR sbstrParamName;
                        hr = sbstrParamName.Append(L"MODULE");
                        if (SUCCEEDED(hr))
                        {
                            CComVariant svarValue;
                            hr = InitVariantFromString(strModule, &svarValue);
                            if (SUCCEEDED(hr))
                            {
                                hr = spxslProc->addParameter(sbstrParamName, svarValue, NULL);
                                if (SUCCEEDED(hr))
                                {
                                    CComVariant varInput((IXMLDOMDocument2*)xml);
                                    hr = spxslProc->put_input(varInput);
                                    if (SUCCEEDED(hr))
                                    {
                                        VARIANT_BOOL fDone;
                                        hr = spxslProc->transform(&fDone);
                                        if (SUCCEEDED(hr))
                                        {
                                            CComVariant svarOutput;
                                            hr = spxslProc->get_output(&svarOutput);
                                            if (SUCCEEDED(hr))
                                            {
                                                CXMLDocument xmlOut;
                                                hr = xmlOut.LoadXML(svarOutput.bstrVal);
                                                if (SUCCEEDED(hr))
                                                {
                                                    CPathString strSchema;
                                                    hr = _BuildSchemaPath(strSchema, strSchema.Size());
                                                    if (SUCCEEDED(hr))
                                                    {
                                                        hr = xmlOut.Save(strSchema);
                                                        if (SUCCEEDED(hr))
                                                        {
                                                            hr = PSRegisterPropertySchema(strSchema);
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    return hr;
}

HRESULT CRegNamespaceModule::UnregisterPropertySchema()
{
    CPathString strSchema;
    HRESULT hr = _BuildSchemaPath(strSchema, strSchema.Size());
    if (SUCCEEDED(hr))
    {
        hr = PSUnregisterPropertySchema(strSchema);
        if (SUCCEEDED(hr))
        {
            DeleteFile(strSchema);
        }
    }
    return hr;
}

// DLL Entry Point
extern "C" BOOL WINAPI DllMain(__in HINSTANCE hInstance, DWORD dwReason, __in void *pvReserved)
{
#ifdef _MERGE_PROXYSTUB
    if (!PrxDllMain(hInstance, dwReason, pvReserved))
        return FALSE;
#endif
    return _AtlModule.DllMain(dwReason, pvReserved); 
}

// Used to determine whether the DLL can be unloaded by OLE
STDAPI DllCanUnloadNow(void)
{
#ifdef _MERGE_PROXYSTUB
    HRESULT hr = PrxDllCanUnloadNow();
    if (hr != S_OK)
        return hr;
#endif
    return _AtlModule.DllCanUnloadNow();
}

// Returns a class factory to create an object of the requested type
STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, __deref_out void **ppv)
{
#ifdef _MERGE_PROXYSTUB
    if (PrxDllGetClassObject(rclsid, riid, ppv) == S_OK)
        return S_OK;
#endif
    return _AtlModule.DllGetClassObject(rclsid, riid, ppv);
}

STDAPI DllRegisterServer(void)
{
    return _AtlModule.DllRegisterServer(FALSE);
}

// DllUnregisterServer - Removes entries from the system registry
STDAPI DllUnregisterServer(void)
{
    return _AtlModule.DllUnregisterServer(FALSE);
}

extern "C" UINT WINAPI ManageRegistrySchema(MSIHANDLE hmsi)
{
    CComMemPtr<WCHAR> spszCustomActionData;
    HRESULT hr = GetMsiProperty(hmsi, IPROPNAME_CUSTOMACTIONDATA, &spszCustomActionData);
    if (SUCCEEDED(hr))
    {
        if (StrCmpIC(spszCustomActionData, L"Install") == 0)
        {
            hr = _AtlModule.RegisterPropertySchema();
        }
        else
        {
            hr = _AtlModule.UnregisterPropertySchema();
        }
        if (SUCCEEDED(hr))
        {
            _AtlModule.NotifyDesktopChange();
        }
    }
    return SCODE_CODE(GetScode(hr));
}
